use linfa::prelude::*;
use linfa_clustering::KMeans;
use ndarray::Array2;

#[derive(Debug, Clone)]
struct PriceTick {
    time: u32,
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    volume: u32,
}

fn price_tick_to_feature_vector(price_tick: &PriceTick) -> Vec<f64> {
    vec![
        price_tick.open,
        price_tick.high,
        price_tick.low,
        price_tick.close,
        price_tick.volume as f64,
    ]
}

fn main() {
    // Define the base K-line to compare against
    let base_kline = vec![
        PriceTick { time: 1, open: 98.0, high: 105.0, low: 97.0, close: 102.0, volume: 100, },
        PriceTick { time: 2, open: 102.0, high: 108.0, low: 100.0, close: 105.0, volume: 120, },
    ];

    // Define additional K-lines for comparison
    let kline1 = vec![
        PriceTick { time: 3, open: 100.0, high: 106.0, low: 98.0, close: 104.0, volume: 110, },
        PriceTick { time: 4, open: 97.0, high: 103.0, low: 95.0, close: 100.0, volume: 105, },
        // Add more PriceTick structs as needed
    ];

    let kline2 = vec![
        PriceTick { time: 3, open: 99.0, high: 105.0, low: 97.0, close: 103.0, volume: 115, },
        PriceTick { time: 4, open: 98.0, high: 104.0, low: 96.0, close: 101.0, volume: 108, },
        // Add more PriceTick structs as needed
    ];

    let kline3 = vec![
        PriceTick { time: 3, open: 101.0, high: 107.0, low: 99.0, close: 105.0, volume: 112, },
        PriceTick { time: 4, open: 96.0, high: 102.0, low: 94.0, close: 99.0, volume: 100, },
        // Add more Pric…
    ];

    // Create a vector of all K-lines for comparison
    let mut klines = vec![kline1, kline2, kline3]; // Add more K-lines as needed
                                                           // Convert K-line data to feature vectors
    let base_features: Vec<Vec<f64>> = base_kline
        .iter()
        .map(|price_tick| price_tick_to_feature_vector(price_tick))
        .collect();

    let kline_features: Vec<Vec<Vec<f64>>> = klines
        .iter()
        .map(|kline| {
            kline
                .iter()
                .map(|price_tick| price_tick_to_feature_vector(price_tick))
                .collect()
        })
        .collect();

    // Convert feature vectors to ndarray Array2 for clustering
    let base_features_array = Array2::from_shape_vec((base_features.len(), 5), base_features).unwrap();

    let kline_features_array = Array2::from_shape_vec(
        (kline_features.len(), 5),
        kline_features.iter().flat_map(|v| v.clone()).collect(),
    )
    .unwrap();
    // Flatten the nested feature vectors into a single vector
    let flattened_features: Vec<Vec<f64>> = kline_features.into_iter().flat_map(|v| v).collect();

    let kline_features_array =
        Array2::from_shape_vec((flattened_features.len(), 5), flattened_features).unwrap();

    // Perform KMeans clustering
    let kmeans = KMeans::params(3) // Specify the number of clusters
    .fit(&DatasetBase::from(base_features_array))
    .unwrap();

    // Perform KMeans clustering
    let kmeans = KMeans::params(3) // Specify the number of clusters
        .fit(&DatasetBase::from(kline_features_array))
        .unwrap();

    // Predict the cluster labels for the K-line features
    let labels = kmeans.predict(&DatasetBase::from(kline_features_array));

    // Find the index of the most similar K-line based on cluster label or distance metric

    println!("Cluster labels for the K-lines: {:?}", labels);
}
